#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
预处理(2017-2022) + 预测所有因子并用其修正就业率(2025-2026, 3/1~7/1 每4天)。
输出三份文件：
  1) pre_2017_2022_grid.csv                # 训练期对齐栅格的多元表
  2) factors_2025_2026_grid.csv            # 未来两年所有因子在4天栅格上的预测
  3) employment_rate_forecast_2025_2026.csv# 未来两年就业率预测（两列：date,就业率（%））
"""

import argparse, re, warnings, os
warnings.filterwarnings("ignore")

import pandas as pd
import numpy as np
from pandas.tseries.offsets import MonthEnd

TRAIN_YEARS = list(range(2017, 2023))
FORE_YEARS  = [2025, 2026]

# ---------------------------- 时间解析（支持季度/年月/常见日期） ----------------------------
def _try_parse_quarter_str(s: pd.Series) -> pd.Series:
    def _to_dt(x):
        if pd.isna(x): return pd.NaT
        xs = str(x).strip()
        zh_map = {"一":"1","二":"2","三":"3","四":"4","Ⅰ":"1","Ⅱ":"2","Ⅲ":"3","Ⅳ":"4"}
        for k,v in zh_map.items(): xs = xs.replace(k, v)
        m = re.match(r"^\s*(\d{4})\s*[-/ ]?\s*[Q第]?\s*([1-4])\s*[季度Q]?\s*$", xs)
        if m:
            y = int(m.group(1)); q = int(m.group(2))
            return pd.Timestamp(year=y, month=q*3, day=1) + MonthEnd(0)  # 季度末
        return pd.NaT
    return s.map(_to_dt)

def _try_parse_yearmonth(s: pd.Series) -> pd.Series:
    def _to_dt(x):
        if pd.isna(x): return pd.NaT
        xs = str(x).strip()
        if re.match(r"^\d{6}$", xs):
            dt = pd.to_datetime(xs, format="%Y%m", errors="coerce")
            return dt + MonthEnd(0) if pd.notna(dt) else pd.NaT
        if re.match(r"^\d{4}[-/.]\d{1,2}$", xs):
            dt = pd.to_datetime(xs, errors="coerce")
            return (dt + MonthEnd(0)) if pd.notna(dt) else pd.NaT
        return pd.NaT
    return s.map(_to_dt)

def _coerce_datetime(df: pd.DataFrame, date_col: str) -> pd.DatetimeIndex:
    ser = df[date_col]
    dt = pd.to_datetime(ser, errors="coerce", infer_datetime_format=True)
    if dt.notna().mean() > 0.8: return pd.DatetimeIndex(dt)
    dt = _try_parse_quarter_str(ser)
    if dt.notna().mean() > 0.8: return pd.DatetimeIndex(dt)
    dt = _try_parse_yearmonth(ser)
    if dt.notna().mean() > 0.8: return pd.DatetimeIndex(dt)
    return pd.DatetimeIndex([pd.NaT]*len(ser))

def guess_and_parse_datetime(df: pd.DataFrame) -> pd.DatetimeIndex:
    for cand in ["date","时间","日期","ds","time"]:
        if cand in df.columns:
            dt = _coerce_datetime(df, cand)
            if dt.notna().any(): return dt
    low = {c.lower(): c for c in df.columns}
    if "year" in low and ("quarter" in low or "season" in low):
        y = pd.to_numeric(df[low["year"]], errors="coerce")
        qn = low["quarter"] if "quarter" in low else low["season"]
        q = pd.to_numeric(df[qn], errors="coerce").clip(1,4)
        dt = pd.to_datetime(dict(year=y, month=(q*3).astype("Int64"), day=1), errors="coerce") + MonthEnd(0)
        return pd.DatetimeIndex(dt)
    if "year" in low and "month" in low:
        y = pd.to_numeric(df[low["year"]], errors="coerce")
        m = pd.to_numeric(df[low["month"]], errors="coerce").clip(1,12)
        dt = pd.to_datetime(dict(year=y, month=m, day=1), errors="coerce") + MonthEnd(0)
        return pd.DatetimeIndex(dt)
    first = df.columns[0]
    dt = _coerce_datetime(df, first)
    if dt.notna().any(): return dt
    raise ValueError("无法识别时间列；可用列名：" + ", ".join(map(str, df.columns)))

# ---------------------------- 读取 employ_rate.csv（决定训练期栅格与 y） ----------------------------
def ensure_percent_scale(s: pd.Series) -> pd.Series:
    s = pd.to_numeric(s, errors="coerce")
    if s.max(skipna=True) <= 1.0: s = s * 100.0
    return s.clip(lower=0, upper=100)

def load_target_grid_and_y(path: str):
    df = pd.read_csv(path)
    date_col = None
    for c in ["date","时间","日期","ds","time"]:
        if c in df.columns: date_col = c; break
    if date_col is None:
        raise ValueError("employ_rate.csv 未找到时间列（date/时间/日期/ds/time）")
    tcol = None
    for c in df.columns:
        if c in {"employment_rate","就业率","就业率（%）"}: tcol = c; break
    if tcol is None:
        for c in df.columns:
            cl = str(c).lower().replace("（%）","").replace("%","").strip()
            if "employment" in cl and "rate" in cl: tcol = c; break
    if tcol is None:
        raise ValueError("employ_rate.csv 未识别到目标列（employment_rate/就业率/就业率（%））")

    df = df.rename(columns={date_col:"date", tcol:"employment_rate"})
    df["date"] = pd.to_datetime(df["date"])
    df = df[(df["date"].dt.year>=2017) & (df["date"].dt.year<=2022)].copy().sort_values("date")
    grid = pd.DatetimeIndex(df["date"])
    y = ensure_percent_scale(df["employment_rate"]).values
    return grid, pd.Series(y, index=grid, name="employment_rate")

# ---------------------------- 训练期：外生变量对齐训练栅格（毕业生人数季度常数，其他插值） ----------------------------
def load_allvars_train_grid(path: str, grid_idx: pd.DatetimeIndex,
                            grad_col: str = None, wanted_cols=None) -> pd.DataFrame:
    raw = pd.read_csv(path)
    dt = guess_and_parse_datetime(raw)
    raw = raw.copy()
    raw["date"] = dt
    raw = raw.dropna(subset=["date"]).sort_values("date").set_index("date")

    num = raw.select_dtypes(include=[np.number]).copy()
    if num.empty:
        raise ValueError("all_vars_yoy.csv 中没有数值列。")

    daily_interp = num.resample("D").mean().interpolate("time").ffill().bfill()
    if wanted_cols is None:
        wanted_cols = list(num.columns)

    def quarter_key(ts):
        q = (ts.month - 1) // 3 + 1
        return pd.Timestamp(year=ts.year, month=q * 3, day=1) + MonthEnd(0)

    X = pd.DataFrame(index=grid_idx)
    for col in wanted_cols:
        if grad_col and col == grad_col:
            # —— 毕业生人数：季度内常数 —— #
            s_q = num[col].copy().sort_index()
            q_map = pd.Series([quarter_key(d) for d in grid_idx], index=grid_idx)
            s_q_filled = s_q.reindex(s_q.index.union(q_map.unique())).sort_index().ffill().bfill()
            X[col] = s_q_filled.reindex(q_map.values).values
        else:
            # —— 连续因子：插值 + ffill/bfill + 季度兜底 —— #
            s = pd.to_numeric(daily_interp[col], errors="coerce")
            s = s.reindex(grid_idx).interpolate("time").ffill().bfill()
            if s.isna().any():
                s_q = num[col].copy().sort_index()
                q_map = pd.Series([quarter_key(d) for d in grid_idx], index=grid_idx)
                s_q_filled = s_q.reindex(s_q.index.union(q_map.unique())).sort_index().ffill().bfill()
                s = s.fillna(pd.Series(s_q_filled.reindex(q_map.values).values, index=grid_idx))
            X[col] = s.values

    if X.isna().sum().sum() > 0:
        print("[WARN] 训练期回归因子仍有缺失，已做 ffill/bfill 兜底：", X.isna().sum()[X.isna().sum()>0].to_dict())
        X = X.ffill().bfill()

    return X

# ---------------------------- 未来：构造 3/1~7/1 每4天的栅格 ----------------------------
def make_4day_grid_for_year(year: int) -> pd.DatetimeIndex:
    d = pd.Timestamp(year=year, month=3, day=1)
    end = pd.Timestamp(year=year, month=7, day=1)
    out = []
    while d <= end:
        out.append(d); d += pd.Timedelta(days=4)
    return pd.DatetimeIndex(out)

def build_future_grid() -> pd.DatetimeIndex:
    return pd.DatetimeIndex(sorted(np.concatenate([make_4day_grid_for_year(y) for y in FORE_YEARS])))

# ---------------------------- 未来：季度常数 & 连续因子 逐列预测 ----------------------------
def forecast_quarterly_const(series_quarter: pd.Series, fut_grid: pd.DatetimeIndex) -> pd.Series:
    s = series_quarter.copy().sort_index().resample("Q").mean().dropna()
    try:
        from prophet import Prophet
        dfp = pd.DataFrame({"ds": s.index, "y": s.values})
        m = Prophet(weekly_seasonality=False, daily_seasonality=False, yearly_seasonality=True)
        m.fit(dfp)
        q_targets = []
        for y in FORE_YEARS:
            q_targets += [pd.Timestamp(f"{y}-03-31"), pd.Timestamp(f"{y}-06-30")]
        futq = pd.DataFrame({"ds": pd.to_datetime(q_targets)})
        yhat = m.predict(futq)[["ds","yhat"]].set_index("ds")["yhat"]
    except Exception:
        idx = s.index.view("i8").values.astype(float)
        vals = s.values.astype(float)
        k, b = np.polyfit(idx, vals, 1)
        q_targets = []
        for y in FORE_YEARS:
            q_targets += [pd.Timestamp(f"{y}-03-31"), pd.Timestamp(f"{y}-06-30")]
        idx2 = pd.to_datetime(q_targets).view("i8").values.astype(float)
        yhat = pd.Series(k*idx2 + b, index=pd.to_datetime(q_targets))

    def quarter_end(ts):
        q = (ts.month-1)//3 + 1
        return pd.Timestamp(year=ts.year, month=q*3, day=1) + MonthEnd(0)
    qe = pd.Series([quarter_end(d) for d in fut_grid], index=fut_grid)
    yhat_full = yhat.reindex(yhat.index.union(qe.unique())).sort_index().ffill().bfill()
    return yhat_full.reindex(qe.values).set_axis(fut_grid)

def forecast_continuous(series_daily: pd.Series, train_grid: pd.DatetimeIndex, fut_grid: pd.DatetimeIndex) -> pd.Series:
    s = pd.to_numeric(series_daily, errors="coerce").asfreq("D").interpolate("time").ffill().bfill()
    s_train = s[(s.index.year>=min(TRAIN_YEARS)) & (s.index.year<=max(TRAIN_YEARS))]
    try:
        from prophet import Prophet
        dfp = pd.DataFrame({"ds": s_train.index, "y": s_train.values})
        m = Prophet(weekly_seasonality=False, daily_seasonality=False, yearly_seasonality=True)
        m.fit(dfp)
        fut = pd.DataFrame({"ds": fut_grid})
        pred = m.predict(fut)[["ds","yhat"]].set_index("ds")["yhat"]
        return pred
    except Exception:
        t0 = (s_train.index - s_train.index[0]).days.values.astype(float)
        y  = s_train.values.astype(float)
        w = 2*np.pi/365.25
        X = np.c_[np.ones_like(t0), t0, np.sin(w*t0), np.cos(w*t0)]
        beta = np.linalg.pinv(X) @ y
        t_pred = (fut_grid - s_train.index[0]).days.values.astype(float)
        Xp = np.c_[np.ones_like(t_pred), t_pred, np.sin(w*t_pred), np.cos(w*t_pred)]
        yhat = Xp @ beta
        return pd.Series(yhat, index=fut_grid, dtype=float)

# ---------------------------- Prophet 多元回归（用因子修正就业率） ----------------------------
def prophet_with_regressors(train_df: pd.DataFrame, fut_df: pd.DataFrame) -> pd.DataFrame:
    from prophet import Prophet
    reg_cols = [c for c in train_df.columns if c not in {"ds","y"}]

    before = len(train_df)
    train_df = train_df.dropna(subset=["y"] + reg_cols)
    dropped = before - len(train_df)
    if dropped > 0:
        print(f"[INFO] 训练集含缺失，已丢弃 {dropped} 行以满足 Prophet 要求。")

    m = Prophet(weekly_seasonality=False, daily_seasonality=False, yearly_seasonality=True)
    for c in reg_cols:
        m.add_regressor(c, standardize=True)
    m.fit(train_df[["ds","y"]+reg_cols])

    fut_df = fut_df.copy()
    for c in reg_cols:
        if fut_df[c].isna().any():
            fut_df[c] = fut_df[c].interpolate().ffill().bfill()

    yhat = m.predict(fut_df[["ds"]+reg_cols])[["ds","yhat"]]
    yhat["yhat"] = yhat["yhat"].clip(0, 100)
    return yhat.rename(columns={"ds":"date","yhat":"就业率（%）"})

# ---------------------------- 主流程 ----------------------------
def main():
    ap = argparse.ArgumentParser(description="预处理(2017-2022) + 因子预测 + Prophet多元预测(2025-2026, 3/1~7/1 每4天)")
    ap.add_argument("--allvars", default="all_vars_yoy.csv", help="外生变量CSV（季度/混频）")
    ap.add_argument("--target",  default="employ_rate.csv",  help="目标就业率CSV")
    ap.add_argument("--out-pre", default="pre_2017_2022_grid.csv")
    ap.add_argument("--out-factors", default="factors_2025_2026_grid.csv")
    ap.add_argument("--out-fore", default="employment_rate_forecast_2025_2026.csv")
    ap.add_argument("--grad-col", default=None, help="毕业生人数列名（季度内常数，不插值）")
    ap.add_argument("--cols", nargs="*", default=None, help="参与回归/预测的列名（至少4个，需在 all_vars_yoy.csv 里）")
    args = ap.parse_args()

    if not os.path.exists(args.allvars) or not os.path.exists(args.target):
        raise SystemExit(f"找不到输入文件：{args.allvars} / {args.target}（可用 --allvars / --target 指定路径）")

    # 1) 训练期：栅格与 y（完全对齐 employ_rate.csv）
    train_grid, y_series = load_target_grid_and_y(args.target)

    # 2) 训练期：外生变量对齐训练栅格
    X_train = load_allvars_train_grid(args.allvars, train_grid, grad_col=args.grad_col, wanted_cols=args.cols)
    pre_df = pd.DataFrame({"date": train_grid, "employment_rate": y_series.values})
    pre_df = pre_df.join(X_train, on="date")
    pre_df.to_csv(args.out_pre, index=False, encoding="utf-8-sig")
    print(f"[PRE] saved {args.out_pre} rows={len(pre_df)}")

    # 3) 未来（2025-2026）：预测所有回归因子在4天栅格
    fut_grid = build_future_grid()
    raw = pd.read_csv(args.allvars)
    dt = guess_and_parse_datetime(raw)
    raw = raw.copy()
    raw["date"] = dt
    raw = raw.dropna(subset=["date"]).sort_values("date").set_index("date")
    num = raw.select_dtypes(include=[np.number]).copy()
    daily_all = num.resample("D").mean().interpolate("time").ffill().bfill()

    factors_future = pd.DataFrame(index=fut_grid)
    use_cols = args.cols if args.cols is not None else list(num.columns)
    for col in use_cols:
        if args.grad_col and col == args.grad_col:
            factors_future[col] = forecast_quarterly_const(num[col], fut_grid)
        else:
            factors_future[col] = forecast_continuous(daily_all[col], train_grid, fut_grid)

    factors_future = factors_future.reset_index().rename(columns={"index":"date"})
    factors_future.to_csv(args.out_factors, index=False, encoding="utf-8-sig")
    print(f"[FACTORS] saved {args.out_factors} rows={len(factors_future)}")

    # 4) Prophet 多元回归预测（仅 2017–2022 训练；在 2025–2026 的4天栅格上预测就业率）
    train_df = pre_df.rename(columns={"date":"ds","employment_rate":"y"}).copy()
    fut_df = factors_future.rename(columns={"date":"ds"}).copy()
    yhat = prophet_with_regressors(train_df, fut_df)
    yhat.to_csv(args.out_fore, index=False, encoding="utf-8-sig")
    print(f"[FORE] saved {args.out_fore} rows={len(yhat)}")

if __name__ == "__main__":
    main()
